#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import unicode_literals, print_function

import codecs
import datetime
import glob
import locale
import os
import re
import subprocess
import sys

DB_USER = 'SYSDBA'
DB_PASSWORD = 'Dameng123'
SYSAUDITOR_PASSWORD = DB_PASSWORD
DB_HOST = '127.0.0.1'
DB_PORT = '5236'
RUN_MINUTES = '3'
WAREHOUSE_COUNT = '100'
TERMINAL_COUNT = '100'
LOAD_WORKERS = '10'
TABLE_TYPE = '10'
PARTITION_COUNT = '100'
TABLESPACE_SIZE_MB = '2048'
TEST_NODE = '1'
TEST_FLAG = '1'
ACID_FLAG = '0'
CKPT_FLAG = '1'
PREHEAT_FLAG = '1'
REDO_LOG_SIZE_MB = '2048'
NUMASTAT_TARGET_LINES = (10, 19, 28)
LOG_FILE_NAME = 'test.log'
LOG_FILE_ENCODING = 'utf-8'
SHELL_ENCODING_FALLBACKS = ('utf-8', 'gbk', 'gb2312', 'big5', 'latin1')
SCRIPT_ENCODING_FALLBACKS = ('utf-8', 'gbk', 'gb2312', 'big5', 'latin1')
SHELL_INTERPRETER = '/bin/sh'
DM_BIN_RELATIVE = os.path.join('..', 'dm8', 'bin')
DM_LOG_RELATIVE = os.path.join('..', 'dm8', 'log')
TPCC_LOG_RELATIVE = os.path.join('log', '01_tpcc')
INIT_SCRIPT_RELATIVE = os.path.join('..', 'init_dm8', '1initDB.sh')
TPCC_LOAD_SCRIPT_RELATIVE = 'tpcc_load.sh'
TPCC_TEST_SCRIPT_RELATIVE = 'tpcc_test.sh'
DISQL_BINARY_NAME = 'disql'
TAIL_LINES = 10
TPCC_RESULT_CONTEXT = 5


class TPCCTestRunner(object):
    def __init__(self):
        self.base_dir = os.path.abspath(os.path.dirname(__file__))
        self.line_ending = '\n'
        self.env = os.environ.copy()
        if 'PYTHONIOENCODING' not in self.env:
            self.env['PYTHONIOENCODING'] = LOG_FILE_ENCODING
        self.output_encoding = locale.getpreferredencoding() or 'utf-8'
        self.console_encoding = getattr(sys.stdout, 'encoding', None) or self.output_encoding
        self.log_path = os.path.join(self.base_dir, LOG_FILE_NAME)
        self.dm_bin_dir = os.path.normpath(os.path.join(self.base_dir, DM_BIN_RELATIVE))
        self.dm_log_dir = os.path.normpath(os.path.join(self.base_dir, DM_LOG_RELATIVE))
        self.tpcc_log_dir = os.path.normpath(os.path.join(self.base_dir, TPCC_LOG_RELATIVE))
        self.config_backup_dir = os.path.join(self.base_dir, 'config_backup')
        self.shell_interpreter = SHELL_INTERPRETER
        self.close_fds = True
        self.disql_path = self._resolve_disql()

    def _resolve_disql(self):
        path = os.path.join(self.dm_bin_dir, DISQL_BINARY_NAME)
        if os.path.exists(path):
            return path
        return path

    def _decode_bytes(self, data):
        if isinstance(data, unicode):
            return data, 'utf-8'
        for encoding in SCRIPT_ENCODING_FALLBACKS:
            try:
                return data.decode(encoding), encoding
            except Exception:
                continue
        return data.decode('utf-8', 'ignore'), 'utf-8'

    def _normalize_newlines(self, text, ensure_trailing_newline):
        normalized = text.replace('\r\n', '\n').replace('\r', '\n')
        if ensure_trailing_newline and not normalized.endswith('\n'):
            normalized += '\n'
        return normalized.replace('\n', self.line_ending)

    def _append_to_log(self, text, ensure_trailing_newline=False):
        if not text:
            return
        fragment = self._normalize_newlines(text, ensure_trailing_newline)
        parent_dir = os.path.dirname(self.log_path)
        if parent_dir and not os.path.exists(parent_dir):
            try:
                os.makedirs(parent_dir)
            except OSError:
                pass
        with codecs.open(self.log_path, 'a', encoding=LOG_FILE_ENCODING) as handler:
            handler.write(fragment)

    def _write_console(self, text):
        if not text:
            return
        try:
            sys.stdout.write(text)
        except UnicodeEncodeError:
            sys.stdout.write(text.encode(self.console_encoding or 'utf-8', 'replace'))
        sys.stdout.flush()

    def log(self, message, write_file=True):
        timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        line = '[%s] %s' % (timestamp, message)
        self._write_console(line + self.line_ending)
        if write_file:
            self._append_to_log(line + '\n')

    def _decode_chunk(self, chunk):
        if isinstance(chunk, unicode):
            return chunk
        for encoding in (self.output_encoding,) + SHELL_ENCODING_FALLBACKS:
            try:
                return chunk.decode(encoding)
            except Exception:
                continue
        return chunk.decode('utf-8', 'ignore')

    def execute_command(self, command, capture_output=False):
        use_shell = isinstance(command, basestring)
        process = subprocess.Popen(
            command,
            shell=use_shell,
            cwd=self.base_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            bufsize=1,
            close_fds=self.close_fds,
            env=self.env
        )
        collected = []
        while True:
            chunk = process.stdout.readline()
            if not chunk:
                if process.poll() is not None:
                    break
                continue
            decoded = self._decode_chunk(chunk)
            self._write_console(decoded)
            if capture_output:
                collected.append(decoded)
        returncode = process.wait()
        if capture_output:
            return returncode, ''.join(collected)
        return returncode, None

    def _find_backup_bytes(self, relative_path):
        if not self.config_backup_dir or not os.path.isdir(self.config_backup_dir):
            return None, None
        base_name = os.path.basename(relative_path)
        pattern = os.path.join(self.config_backup_dir, '%s.*' % base_name)
        candidates = [path for path in glob.glob(pattern) if os.path.isfile(path)]
        if not candidates:
            return None, None
        latest = max(candidates, key=os.path.getmtime)
        with open(latest, 'rb') as handler:
            return handler.read(), latest

    def _load_script_text(self, relative_path):
        absolute = os.path.normpath(os.path.join(self.base_dir, relative_path))
        raw = None
        if os.path.exists(absolute):
            with open(absolute, 'rb') as handler:
                raw = handler.read()
        else:
            self.log('脚本不存在: %s' % relative_path, write_file=False)
        if not raw:
            backup_bytes, backup_path = self._find_backup_bytes(relative_path)
            if backup_bytes:
                raw = backup_bytes
                self.log('脚本内容为空，使用备份: %s' % os.path.basename(backup_path), write_file=False)
            else:
                self.log('无法获取脚本内容: %s' % relative_path)
                return None, None, None, absolute
        newline = '\r\n' if b'\r\n' in raw else '\n'
        text, encoding = self._decode_bytes(raw)
        return text, encoding, newline, absolute

    def _write_script_text(self, path, text, encoding, newline):
        if text is None:
            return
        normalized = text.replace('\r\n', '\n').replace('\r', '\n')
        normalized = normalized.replace('\n', newline)
        with codecs.open(path, 'w', encoding=encoding or 'utf-8') as handler:
            handler.write(normalized)

    def _set_shell_variable(self, content, name, value, quoted):
        if content is None:
            return content
        pattern = re.compile(r'(^\s*%s\s*=\s*)([^\r\n]*?)\s*(#.*)?$' % re.escape(name), re.MULTILINE)

        def _replacer(match):
            prefix = match.group(1)
            comment = match.group(3) or ''
            new_value = value
            if quoted:
                new_value = '"%s"' % value
            line = prefix + new_value
            if comment:
                if not comment.startswith(' '):
                    line += ' '
                line += comment
            return line

        new_content, count = pattern.subn(_replacer, content)
        if count == 0:
            self.log('未在脚本中找到变量 %s' % name, write_file=False)
        return new_content

    def _replace_inline_assignment(self, content, key, value, silent=False):
        if content is None:
            return content
        pattern = re.compile(r'(%s=)([^ \t\r\n]+)' % re.escape(key))
        new_content, count = pattern.subn(r'\1%s' % value, content)
        if count == 0 and not silent:
            self.log('未在脚本中找到参数 %s' % key, write_file=False)
        return new_content

    def _ensure_dminit_parameter(self, content, key, value):
        if content is None:
            return content
        pattern = re.compile(r'(^\s*\./dminit[^\r\n]*)', re.MULTILINE)

        def _inject(match):
            line = match.group(0)
            if '%s=' % key in line:
                return line
            return line.rstrip() + ' %s=%s' % (key, value)

        new_content, count = pattern.subn(_inject, content, count=1)
        if count == 0:
            self.log('未在脚本中找到 dminit 命令，无法追加参数 %s' % key, write_file=False)
            return content
        return new_content


    def _safeguard_kill_command(self, content):
        if content is None:
            return content
        frag = "kill -9 $(ps -ef|grep ${dm_home}/bin |gawk '$0 !~/grep/ {print $2}' |tr -s '\n' ' ')"
        if frag not in content:
            return content
        replacement = (
            "pids=$(ps -ef|grep ${dm_home}/bin |gawk '$0 !~/grep/ {print $2}' |tr -s '\\n' ' ')" + '\n'
            "if [ -n \"$pids\" ]; then" + '\n'
            "    kill -9 $pids" + '\n'
            "fi"
        )
        return content.replace(frag, replacement)

    def _replace_disql_invocation(self, content):
        if content is None:
            return content
        pattern = r'([^\s"]*disql[^\s"]*\s+)[^/\s]+/[^@]+@[^:]+:\d+'
        replacement = r'\1%s/%s@%s:%s' % (DB_USER, DB_PASSWORD, DB_HOST, DB_PORT)
        new_content, count = re.subn(pattern, replacement, content)
        if count == 0:
            legacy = './disql SYSDBA/hust4400_HZQ@localhost:25236'
            if legacy in content:
                new_content = content.replace(legacy, './disql %s/%s@%s:%s' % (
                    DB_USER, DB_PASSWORD, DB_HOST, DB_PORT))
            else:
                self.log('未能更新 disql 连接串', write_file=False)
                new_content = content
        return new_content

    def _set_file_executable(self, path):
        try:
            mode = os.stat(path).st_mode
            os.chmod(path, mode | 0o111)
        except Exception as exc:
            self.log('设置执行权限失败: %s (%s)' % (path, exc), write_file=False)

    def update_script_parameters(self):
        self._update_init_script()
        self._update_tpcc_load_script()
        self._update_tpcc_test_script()

    def _update_init_script(self):
        content, encoding, newline, path = self._load_script_text(INIT_SCRIPT_RELATIVE)
        if content is None:
            return
        content = self._set_shell_variable(content, 'redo_size', REDO_LOG_SIZE_MB, quoted=True)
        content = self._replace_inline_assignment(content, 'SYSDBA_PWD', DB_PASSWORD)
        content = self._replace_inline_assignment(content, 'SYSAUDITOR_PWD', SYSAUDITOR_PASSWORD)
        content = self._replace_inline_assignment(content, 'port_num', DB_PORT, silent=True)
        if 'port_num=' not in content:
            content = self._ensure_dminit_parameter(content, 'port_num', DB_PORT)
        content = self._replace_disql_invocation(content)
        content = self._safeguard_kill_command(content)
        self._write_script_text(path, content, encoding, newline)
        self._set_file_executable(path)

    def _update_tpcc_load_script(self):
        content, encoding, newline, path = self._load_script_text(TPCC_LOAD_SCRIPT_RELATIVE)
        if content is None:
            return
        content = self._set_shell_variable(content, 'DB_HOST1', DB_HOST, quoted=True)
        content = self._set_shell_variable(content, 'DB_PORT1', DB_PORT, quoted=True)
        content = self._set_shell_variable(content, 'SYSDBA_PWD', DB_PASSWORD, quoted=True)
        content = self._set_shell_variable(content, 'TABLE_TYPE', TABLE_TYPE, quoted=True)
        content = self._set_shell_variable(content, 'TBS_SIZE', TABLESPACE_SIZE_MB, quoted=True)
        content = self._set_shell_variable(content, 'WAREHOUSES', WAREHOUSE_COUNT, quoted=True)
        content = self._set_shell_variable(content, 'PARTITION', PARTITION_COUNT, quoted=True)
        content = self._set_shell_variable(content, 'LOADWORKERS', LOAD_WORKERS, quoted=True)
        self._write_script_text(path, content, encoding, newline)
        self._set_file_executable(path)

    def _update_tpcc_test_script(self):
        content, encoding, newline, path = self._load_script_text(TPCC_TEST_SCRIPT_RELATIVE)
        if content is None:
            return
        content = self._set_shell_variable(content, 'DB_HOST1', DB_HOST, quoted=True)
        content = self._set_shell_variable(content, 'DB_PORT1', DB_PORT, quoted=True)
        content = self._set_shell_variable(content, 'SYSDBA_PWD', DB_PASSWORD, quoted=True)
        content = self._set_shell_variable(content, 'RUNMINS', RUN_MINUTES, quoted=True)
        content = self._set_shell_variable(content, 'WAREHOUSES', WAREHOUSE_COUNT, quoted=True)
        content = self._set_shell_variable(content, 'TERMINALS', TERMINAL_COUNT, quoted=True)
        content = self._set_shell_variable(content, 'TEST_NODE', TEST_NODE, quoted=True)
        content = self._set_shell_variable(content, 'TEST_FLAG', TEST_FLAG, quoted=True)
        content = self._set_shell_variable(content, 'ACID_FLAG', ACID_FLAG, quoted=True)
        content = self._set_shell_variable(content, 'CKPT_FLAG', CKPT_FLAG, quoted=True)
        content = self._set_shell_variable(content, 'PREHEAT_FALG', PREHEAT_FLAG, quoted=True)
        self._write_script_text(path, content, encoding, newline)
        self._set_file_executable(path)

    def run_shell_script(self, relative_path, description):
        path = os.path.normpath(os.path.join(self.base_dir, relative_path))
        if not os.path.exists(path):
            self.log('未找到脚本: %s' % path)
            return False
        self.log('开始执行%s' % description)
        command = [self.shell_interpreter, path]
        returncode, _ = self.execute_command(command)
        if returncode != 0:
            self.log('%s 执行失败 (退出码 %s)' % (description, returncode))
            return False
        self.log('%s 执行完成' % description)
        return True

    def collect_numa_stats(self):
        self.log('收集 NUMA 统计', write_file=False)
        username = os.environ.get('USER') or os.environ.get('LOGNAME') or ''
        command = ['pgrep', '-u', username, 'dmserver'] if username else ['pgrep', 'dmserver']
        code, output = self.execute_command(command, capture_output=True)
        if code != 0 or not output.strip():
            self.log('未找到 dmserver 进程，跳过 NUMA 统计', write_file=True)
            return
        pid = output.strip().splitlines()[-1].strip()
        if not pid:
            self.log('无法解析 dmserver PID，跳过 NUMA 统计', write_file=True)
            return
        code, stats = self.execute_command(['numastat', '-p', pid], capture_output=True)
        if code != 0 or not stats:
            self.log('numastat 执行失败，跳过 NUMA 统计', write_file=True)
            return
        lines = stats.splitlines()
        selected = []
        for index in NUMASTAT_TARGET_LINES:
            position = index - 1
            if position < len(lines):
                selected.append(lines[position])
        if selected:
            self._append_to_log('\n'.join(selected) + '\n')
            self.log('NUMA 统计已写入日志', write_file=False)
        else:
            self.log('numastat 输出不足，未写入日志', write_file=True)

    def _extract_disql_value(self, output):
        if not output:
            return ''
        filtered = []
        for line in output.splitlines():
            cleaned = line.strip()
            if not cleaned:
                continue
            lower = cleaned.lower()
            if lower.startswith('sql>'):
                continue
            if 'row' in lower and 'fetched' in lower:
                continue
            filtered.append(cleaned)
        if not filtered:
            return ''
        return filtered[-1]

    def collect_db_params(self):
        self.log('收集数据库参数', write_file=False)
        if not os.path.exists(self.disql_path):
            self.log('disql 可执行文件不存在: %s' % self.disql_path, write_file=True)
            return
        connection = '%s/%s@%s:%s' % (DB_USER, DB_PASSWORD, DB_HOST, DB_PORT)
        tasks = [
            ('ENABLE_FREQROOTS', "select para_value from v$dm_ini where para_name='ENABLE_FREQROOTS'"),
            ('parttions', 'select count(*) from dba_tab_partitions')
        ]
        for label, sql in tasks:
            command = [self.disql_path, connection, '-e', sql]
            code, output = self.execute_command(command, capture_output=True)
            if code != 0:
                self.log('执行 SQL 失败: %s' % sql, write_file=True)
                continue
            value = self._extract_disql_value(output)
            if value:
                self._append_to_log('%s:%s\n' % (label, value))
            else:
                self.log('未获取到 %s 的返回值' % label, write_file=True)
        self.log('数据库参数收集完成', write_file=False)

    def _read_text_file(self, path):
        with open(path, 'rb') as handler:
            raw = handler.read()
        return self._decode_bytes(raw)

    def collect_transaction_info(self):
        self.log('收集事务提交信息', write_file=False)
        month_tag = datetime.datetime.now().strftime('%Y%m')
        log_path = os.path.join(self.dm_log_dir, 'dm_DMSERVER_%s.log' % month_tag)
        if not os.path.exists(log_path):
            self.log('未找到 DM 日志文件: %s' % log_path, write_file=True)
            return
        content, _ = self._read_text_file(log_path)
        lines = content.splitlines()
        tail = lines[-TAIL_LINES:]
        matched = []
        for line in tail:
            if 'trx4_cmt' in line:
                matched.append(line)
        if matched:
            self._append_to_log('\n'.join(matched) + '\n')
            self.log('事务信息已写入日志', write_file=False)
        else:
            self.log('最近日志未包含 trx4_cmt', write_file=True)

    def collect_tpcc_results(self):
        self.log('收集 TPCC 结果', write_file=False)
        if not os.path.isdir(self.tpcc_log_dir):
            self.log('TPCC 日志目录不存在: %s' % self.tpcc_log_dir, write_file=True)
            return
        files = [path for path in glob.glob(os.path.join(self.tpcc_log_dir, '*')) if os.path.isfile(path)]
        if not files:
            self.log('TPCC 日志目录为空', write_file=True)
            return
        latest = max(files, key=os.path.getmtime)
        content, _ = self._read_text_file(latest)
        lines = content.splitlines()
        pattern = re.compile(r'Measured tpmC \(NewOrders\)')
        block = []
        for index, line in enumerate(lines):
            if pattern.search(line):
                block = lines[index:index + TPCC_RESULT_CONTEXT]
                break
        if block:
            self._append_to_log('\n'.join(block) + '\n')
            self.log('TPCC 结果已写入日志: %s' % os.path.basename(latest), write_file=False)
        else:
            self.log('在 %s 中未找到 tpmC 结果' % os.path.basename(latest), write_file=True)

    def run(self):
        self.log('========== TPCC 测试开始 ==========')
        self.update_script_parameters()
        steps = [
            (INIT_SCRIPT_RELATIVE, '数据库初始化'),
            (TPCC_LOAD_SCRIPT_RELATIVE, 'TPCC 数据装载'),
            (TPCC_TEST_SCRIPT_RELATIVE, 'TPCC 性能测试')
        ]
        for script_path, description in steps:
            if not self.run_shell_script(script_path, description):
                self.log('流程终止: %s 执行失败' % description)
                return 1
        self.log('开始收集运行指标', write_file=False)
        self.collect_numa_stats()
        self.collect_db_params()
        self.collect_transaction_info()
        self.collect_tpcc_results()
        self.log('========== TPCC 测试结束 ==========')
        self.log('测试日志输出位置: %s' % self.log_path)
        return 0


def main():
    runner = TPCCTestRunner()
    try:
        exit_code = runner.run()
    except KeyboardInterrupt:
        runner.log('检测到用户终止', write_file=True)
        exit_code = 1
    except Exception as exc:
        runner.log('执行过程中发生错误: %s' % exc, write_file=True)
        exit_code = 1
    sys.exit(exit_code)


if __name__ == '__main__':
    main()
